package index

import (
	"bytes"
	"gitee.com/Cauchy_AQ/gokv/data"
	goart "github.com/plar/go-adaptive-radix-tree"
	"sort"
	"sync"
)

// AdaptiveRadixTree 自适应基数树
type AdaptiveRadixTree struct {
	tree goart.Tree
	lock *sync.RWMutex
}

// NewART 初始化自适应基数树索引
func NewART() *AdaptiveRadixTree {
	return &AdaptiveRadixTree{
		tree: goart.New(),
		lock: new(sync.RWMutex),
	}
}

// Put 向索引中存储 key 对应的数据位置信息
func (a *AdaptiveRadixTree) Put(key []byte, pos *data.LogRecordPos) *data.LogRecordPos {
	a.lock.Lock()
	defer a.lock.Unlock()
	// 不允许往 ART 插入空的 key 值
	if len(key) == 0 {
		return nil
	}
	oldValue, _ := a.tree.Insert(key, pos)
	if oldValue == nil {
		return nil
	}
	return oldValue.(*data.LogRecordPos)
}

// Get 根据 key 取出对于索引位置信息
func (a *AdaptiveRadixTree) Get(key []byte) *data.LogRecordPos {
	a.lock.RLock()
	defer a.lock.RUnlock()
	val, ok := a.tree.Search(key)
	if !ok {
		return nil
	}
	return val.(*data.LogRecordPos)
}

// Delete 根据 key 删除对于索引位置信息
func (a *AdaptiveRadixTree) Delete(key []byte) (*data.LogRecordPos, bool) {
	a.lock.Lock()
	defer a.lock.Unlock()
	value, ok := a.tree.Delete(key)
	if !ok {
		return nil, ok
	}
	return value.(*data.LogRecordPos), ok
}

// Size 获取数据数量大小
func (a *AdaptiveRadixTree) Size() int {
	a.lock.RLock()
	defer a.lock.RUnlock()

	return a.tree.Size()
}

// Close 关闭索引
func (a *AdaptiveRadixTree) Close() error {
	return nil
}

// Iterator 索引迭代器
func (a *AdaptiveRadixTree) Iterator(reverse bool) Iterator {
	a.lock.RLock()
	defer a.lock.RUnlock()
	return newARTIterator(a.tree, reverse)
}

// artIterator AdaptiveRadixTree 索引迭代器
type artIterator struct {
	currentIndex int     // 当前遍历的下标
	reverse      bool    // 是否反向遍历
	values       []*Item // 存储遍历出来的 key 和索引位置信息 data.LogRecordPos
}

// newARTIterator 新建 AdaptiveRadixTree 索引迭代器
func newARTIterator(tree goart.Tree, reverse bool) *artIterator {
	var idx int
	if reverse {
		idx = tree.Size() - 1
	}
	values := make([]*Item, tree.Size())
	saveValues := func(node goart.Node) bool {
		item := &Item{
			key: node.Key(),
			pos: node.Value().(*data.LogRecordPos),
		}
		values[idx] = item
		if reverse {
			idx--
		} else {
			idx++
		}
		return true
	}

	tree.ForEach(saveValues)

	return &artIterator{
		currentIndex: 0,
		reverse:      reverse,
		values:       values,
	}
}

// Rewind 重新返回迭代起点
func (a *artIterator) Rewind() {
	a.currentIndex = 0
}

// Seek 根据传入 key 找到第一个大于（或小于）等于的目标 key，根据这个 key 开始遍历
func (a *artIterator) Seek(key []byte) {
	if a.reverse {
		// 反序，找到第一个小于等于 key 的位置
		a.currentIndex = sort.Search(len(a.values), func(i int) bool {
			return bytes.Compare(a.values[i].key, key) <= 0
		})
	} else {
		// 正序，找到第一个大于等于 key 的位置
		a.currentIndex = sort.Search(len(a.values), func(i int) bool {
			return bytes.Compare(a.values[i].key, key) >= 0
		})
	}
}

// Next 跳转到下一个 key
func (a *artIterator) Next() {
	a.currentIndex += 1
}

// Valid 是否有效，判断是否遍历完了 key
func (a *artIterator) Valid() bool {
	return a.currentIndex < len(a.values)
}

// Key 获取迭代器当前位置的 key
func (a *artIterator) Key() []byte {
	return a.values[a.currentIndex].key
}

// Value 获取当前迭代位置的索引位置信息
func (a *artIterator) Value() *data.LogRecordPos {
	return a.values[a.currentIndex].pos
}

// Close 关闭迭代器
func (a *artIterator) Close() {
	a.values = nil
}
